"""
图像抠图模块

用于自动检测和去除图像背景，保留主体部分
使用AI模型进行精确的前景背景分离
"""

import cv2
import numpy as np
from rembg import remove
from PIL import Image
import io


class ImageMatting:
    """
    图像抠图处理类
    
    功能：
    - 使用AI模型自动去除背景
    - 支持多种图像格式
    - 生成透明背景或指定背景色的图像
    """
    
    def __init__(self):
        """初始化抠图处理器"""
        pass
    
    def remove_background(self, image_path: str) -> np.ndarray:
        """
        去除图像背景
        
        参数:
            image_path: 输入图像路径
            
        返回:
            numpy.ndarray: 去除背景后的RGBA图像数组
            
        使用场景:
            processor = ImageMatting()
            result = processor.remove_background("photo.jpg")
        """
        try:
            # 读取输入图像
            with open(image_path, 'rb') as f:
                input_data = f.read()
            
            # 使用rembg去除背景
            output_data = remove(input_data)
            
            # 转换为PIL图像
            image = Image.open(io.BytesIO(output_data))
            
            # 转换为numpy数组
            image_array = np.array(image)
            
            return image_array
            
        except Exception as e:
            print(f"背景去除失败: {e}")
            # 如果AI抠图失败，使用传统方法作为备选
            return self._fallback_remove_background(image_path)
    
    def _fallback_remove_background(self, image_path: str) -> np.ndarray:
        """
        备选背景去除方法（传统图像处理）
        
        参数:
            image_path: 输入图像路径
            
        返回:
            numpy.ndarray: 处理后的图像数组
        """
        # 读取图像
        image = cv2.imread(image_path)
        
        # 转换色彩空间
        hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        
        # 创建背景掩码（假设背景为白色或浅色）
        lower_white = np.array([0, 0, 200])
        upper_white = np.array([180, 30, 255])
        mask = cv2.inRange(hsv, lower_white, upper_white)
        
        # 形态学操作优化掩码
        kernel = np.ones((3, 3), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
        
        # 创建RGBA图像
        image_rgba = cv2.cvtColor(image, cv2.COLOR_BGR2RGBA)
        image_rgba[mask > 0] = [255, 255, 255, 0]  # 背景设为透明
        
        return image_rgba
    
    def add_white_background(self, image_array: np.ndarray) -> np.ndarray:
        """
        为透明背景图像添加白色背景
        
        参数:
            image_array: 带透明通道的图像数组
            
        返回:
            numpy.ndarray: 添加白色背景后的RGB图像
            
        使用场景:
            result = processor.add_white_background(transparent_image)
        """
        if image_array.shape[2] == 4:  # RGBA图像
            # 创建白色背景
            height, width = image_array.shape[:2]
            white_bg = np.ones((height, width, 3), dtype=np.uint8) * 255
            
            # 提取alpha通道
            alpha = image_array[:, :, 3] / 255.0
            
            # 混合前景和背景
            for c in range(3):
                white_bg[:, :, c] = (
                    alpha * image_array[:, :, c] + 
                    (1 - alpha) * white_bg[:, :, c]
                )
            
            return white_bg
        else:
            # 如果不是RGBA，直接返回RGB
            return image_array[:, :, :3] 